// pmttrans.c

#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <netinet/in.h>
#include <pthread.h>
#include <signal.h>
#include <sys/times.h>
#include <sys/select.h>
#include <unistd.h>

#define PMT_SERVER_SOCKET "/tmp/.listen.camd.socket"
#define PMT_SERVER_SOCKET_BACKUP "/tmp/listen.camd.socket"
#define PMT_LISTEN_SOCKET "/tmp/camd.socket"

typedef unsigned char	bool;
#define true			1
#define false			0

static int pmt_socket_backuped = 0;

static unsigned char pmt_buffer[2048];
static int pmt_len = 0;

static pthread_mutex_t pmt_distribute_lock;
static pthread_cond_t new_pmt_arrived;

static pthread_mutex_t connections_table_lock;
#define SIZE_OF_CONNECTIONS_TABLE	10
static int client_fd[SIZE_OF_CONNECTIONS_TABLE];

static int newpmtfd = -1;

static int delay = 0;

void utils_dump(char *label, unsigned char *data, int len)
{
	int i;
	printf(label);
	for(i = 0; i < len; i++) {
		printf("%02X ", data[i]);
        }
	printf(" [len=%02X]\n", len);
}

int getPMTSize(unsigned char *pmt)
{
	int wp=0;
	if ( pmt[3] & 0x80 )
	{
		int i=0;
		int lenbytes = pmt[3] & ~0x80;
		while(i < lenbytes)
			wp = (wp << 8) | pmt[4 + i++];
		wp+=4;
		wp+=lenbytes;
	}
	else
	{
		wp = pmt[3];
		wp+=4;
	}
	return wp;
}

int is_pmt_valid(unsigned char *pmt, int pmt_size)
{
	if (pmt_size < 31)
		return 0;
	if (memcmp(pmt, "\x9f\x80\x32", 3) == 0 && getPMTSize(pmt) <= pmt_size)
		return 1;
	return 0;
}

struct PMT_ENTRY
{
	unsigned char stream_type;
	unsigned char pid[2];
	unsigned char ca_descriptor_len[2];
	unsigned char *ca_descriptor;
	bool processed;
};

int parse_ca_descriptor(const unsigned char *ca_descriptor, int len, unsigned char **descriptor)
{
	const unsigned char *p = NULL;
	int i = 0;
	while (i < len)
	{
		if (ca_descriptor[i] == 0x09)
		{
			p = ca_descriptor + i;
			break;
		}
		i++;
	}

	int count = 0;
	if (p != NULL)
	{
		while (p < ca_descriptor + len - 1 && p + p[1] + 2 <= ca_descriptor + len)
		{
			descriptor[count] = (unsigned char*)p;
			count++;
			p = p + p[1] + 2;
		}
	}
	return count;
}

bool compare_ca_descriptor(const unsigned char *ca_descriptor1, int len1, const unsigned char *ca_descriptor2, int len2)
{
	if (len1 != len2)
		return false;
	unsigned char *descriptor1[20], *descriptor2[20];
	int count1, count2;
	count1 = parse_ca_descriptor(ca_descriptor1, len1, descriptor1);
	count2 = parse_ca_descriptor(ca_descriptor2, len2, descriptor2);
	if (count1 != count2)
		return false;
	int i, j;
	for (i = 0; i < count1; i++)
	{
		bool found = false;
		for (j = 0; j < count2; j++)
		{
			if (descriptor1[i][1] == descriptor2[j][1] &&
				memcmp(descriptor1[i], descriptor2[j], descriptor1[i][1] + 2) == 0)
			{
				found = true;
				break;
			}
		}
		if (!found)
			return false;
	}
	return true;
}

int getPMTLenBytes(unsigned char *pmt)
{
	int lenbytes = 1;
	if ( pmt[3] & 0x80 )
		lenbytes = (pmt[3] & ~0x80) + 1;
	return lenbytes;
}

int getPMTBaseInfoLen(unsigned char *pmt)
{
	int lenbytes = getPMTLenBytes(pmt);
	int len = (pmt[7 + lenbytes] << 8) + pmt[7 + lenbytes + 1];
	return (7 + lenbytes + 2 + len);
}

int copy_pmt_base_info(unsigned char *new_pmt, unsigned char *original_pmt)
{
	int lenbytes = getPMTLenBytes(original_pmt);
	int len = getPMTSize(original_pmt);
	int base_info_len = getPMTBaseInfoLen(original_pmt);
	memcpy(new_pmt, original_pmt, 3);
	new_pmt[3] = 0x82;
	new_pmt[4] = (len >> 8) & 0xFF;
	new_pmt[5] = len & 0xFF;
	memcpy(new_pmt + 6, original_pmt + 3 + lenbytes, base_info_len - 3 - lenbytes);
	return base_info_len - lenbytes + 3;
}

int convert_single_pmt(unsigned char *original_pmt, int pmt_size, unsigned char *new_pmt)
{
	if (is_pmt_valid(original_pmt, pmt_size) == 0)
	{
		memcpy(new_pmt, original_pmt, pmt_size);
		return pmt_size;
	}

	struct PMT_ENTRY pmt_entry[(pmt_size - 31 + 4) / 5];
	int pmt_entry_count = 0;

	int base_info_len = getPMTBaseInfoLen(original_pmt);
	int wp = base_info_len;
	int len;

	while (wp < pmt_size)
	{
		memcpy(&(pmt_entry[pmt_entry_count]), original_pmt + wp, 5);
		len = (pmt_entry[pmt_entry_count].ca_descriptor_len[0] << 8) + pmt_entry[pmt_entry_count].ca_descriptor_len[1];
		wp += 5;
		if (len == 0)
			pmt_entry[pmt_entry_count].ca_descriptor = NULL;
		else
		{
			pmt_entry[pmt_entry_count].ca_descriptor = original_pmt + wp;
			wp += len;
		}
		pmt_entry[pmt_entry_count].processed = false;
		pmt_entry_count++;
	}

	int i, j;
	int bp = 0, np = 0;
	int processed_count = 0;
	for (i = 0; i < pmt_entry_count; i++)
	{
		if (pmt_entry[i].processed)
			continue;

		np = copy_pmt_base_info(new_pmt + bp, original_pmt);

		j = i;
		do
		{
			memcpy(new_pmt + bp + np, &(pmt_entry[j]), 5);
			np += 5;
			len = (pmt_entry[j].ca_descriptor_len[0] << 8) + pmt_entry[j].ca_descriptor_len[1];
			if (len > 0)
			{
				memcpy(new_pmt + bp + np, pmt_entry[j].ca_descriptor, len);
				np += len;
			}
			pmt_entry[j].processed = true;
			processed_count++;

			do
			{
				j++;
			} while (j < pmt_entry_count &&
				!compare_ca_descriptor(
						pmt_entry[i].ca_descriptor,
						(pmt_entry[i].ca_descriptor_len[0] << 8) + pmt_entry[i].ca_descriptor_len[1],
						pmt_entry[j].ca_descriptor,
						(pmt_entry[j].ca_descriptor_len[0] << 8) + pmt_entry[j].ca_descriptor_len[1]
					)
				);
		} while (j < pmt_entry_count);

		len = np - 6;
		new_pmt[bp + 4] = (len >> 8) & 0xFF;
		new_pmt[bp + 5] = len & 0xFF;

		if (bp == 0)
		{
			if (processed_count == pmt_entry_count)
				new_pmt[bp + 6] = 3;
			else
				new_pmt[bp + 6] = 1;
		}
		else
		{
			if (processed_count == pmt_entry_count)
				new_pmt[bp + 6] = 2;
			else
				new_pmt[bp + 6] = 0;
		}
		bp += np;
	}

	return bp;
}

int convert_pmt(unsigned char *original_pmt, int pmt_size, unsigned char *new_pmt)
{
	int pos = 0, new_pmt_len = 0;
	int single_pmt_len;
	while ((pos + 31) < pmt_size)
	{
		single_pmt_len = getPMTSize(original_pmt + pos);
		if (pos + single_pmt_len <= pmt_size)
			new_pmt_len += convert_single_pmt(original_pmt + pos, single_pmt_len, new_pmt + new_pmt_len);
		pos += single_pmt_len;
	}

	pos = 0;
	while (pos < new_pmt_len)
	{
		single_pmt_len = getPMTSize(new_pmt + pos);
		if (pos == 0 && pos + single_pmt_len < new_pmt_len)
			new_pmt[pos + 6] = 1;
		else if (pos > 0 && pos + single_pmt_len < new_pmt_len)
			new_pmt[pos + 6] = 0;
		else if (pos > 0 && pos + single_pmt_len >= new_pmt_len)
			new_pmt[pos + 6] = 2;
		pos += single_pmt_len;
	}

	return new_pmt_len;
}

void add_to_connections_table(int clientfd)
{
	int i;
	pthread_mutex_lock(&connections_table_lock);
	for (i = 0; i < SIZE_OF_CONNECTIONS_TABLE; i++)
	{
		if (client_fd[i] == -1)
		{
			client_fd[i] = clientfd;
			break;
		}
	}
	pthread_mutex_unlock(&connections_table_lock);
}

void *new_pmt_server_proc(void *pdata)
{
	struct sockaddr_un saddr;

	printf("Trying start new PMT listen socket...\n");

	/* socket init */
	saddr.sun_family = AF_UNIX;
	strncpy(saddr.sun_path, PMT_SERVER_SOCKET, 107);
	saddr.sun_path[107] = '\0';
	unlink(saddr.sun_path);
	if((newpmtfd = socket(AF_UNIX, SOCK_STREAM, 0)) < 0) {
		printf("socket error (errno=%d)\n", errno);
		perror("");
		exit(1);
	}
	if (bind(newpmtfd, (struct sockaddr *)&saddr,
		sizeof(saddr.sun_family)+strlen(saddr.sun_path) ) < 0) {
		printf("socket bind error (errno=%d)\n", errno);
		perror("");
		exit(2);
	}
	if (listen(newpmtfd, 5) < 0) {
		printf("socket listen error (errno=%d)\n", errno);
		perror("");
		exit(3);
	}

	pthread_mutex_init(&connections_table_lock, NULL);

	while(1)
	{
		struct sockaddr_in client_addr;
		socklen_t length = sizeof(client_addr);

		int clientfd = accept(newpmtfd,(struct sockaddr*)&client_addr,&length);
		if (clientfd < 0)
		{
			printf("PMT server accept error (errno=%d)\n", errno);
			perror("");
			break;
		}

		pthread_mutex_lock(&pmt_distribute_lock);
		if (send(clientfd, pmt_buffer, pmt_len, 0) == pmt_len)
		{
			add_to_connections_table(clientfd);
			printf("A client connected.\n");
		}
		pthread_mutex_unlock(&pmt_distribute_lock);
	}

	pthread_exit(NULL);
}

void *pmt_distribute_proc(void *pdata)
{
	int i, len;

	pthread_mutex_lock(&pmt_distribute_lock);

	while(pthread_cond_wait(&new_pmt_arrived, &pmt_distribute_lock) == 0)
	{
		if (delay > 0)
		{
			printf("delaying, %d microseconds\n", delay);
			usleep(delay * 1000);
		}
		
		pthread_mutex_lock(&connections_table_lock);
		for (i = 0; i < SIZE_OF_CONNECTIONS_TABLE; i++)
		{
			if (client_fd[i] != -1)
			{
				printf("sending, client fd is %d\n", client_fd[i]);
				len = send(client_fd[i], pmt_buffer, pmt_len, MSG_NOSIGNAL);
				printf("sended, return value is %d\n", len);
				if (len != pmt_len)
				{
					printf("send error, client fd is %d, return value is %d, errno=%d\n", client_fd[i], len, errno);
					perror("");
					close(client_fd[i]);
					client_fd[i] = -1;
				}
			}
		}
		pthread_mutex_unlock(&connections_table_lock);
	}

	pthread_mutex_unlock(&pmt_distribute_lock);

	pthread_exit(NULL);
}

void init_daemon(void)
{
	int pid;
	int fd, fdtablesize;

	/* 忽略终端 I/O信号，STOP信号 */
	signal(SIGTTOU, SIG_IGN);
	signal(SIGTTIN, SIG_IGN);
	signal(SIGTSTP, SIG_IGN);
	signal(SIGHUP, SIG_IGN);

	/*忽略网络断开信号*/
	signal(SIGPIPE, SIG_IGN);

	if((pid = fork()))
		exit(0);		//是父进程，结束父进程
	else if(pid < 0)
		exit(1);		//fork失败，退出

	//是第一子进程，后台继续执行
	setsid();		//第一子进程成为新的会话组长和进程组长并与控制终端分离

	if((pid = fork()))
		exit(0);		//是第一子进程，结束第一子进程
	else if(pid < 0)
		exit(1);		//fork失败，退出
	//是第二子进程，继续
	//第二子进程不再是会话组长

	for (fd = 0, fdtablesize = getdtablesize(); fd < fdtablesize; fd++)
		close (fd);	//关闭打开的文件描述符

	chdir("/tmp");	//改变工作目录到/tmp
	umask(0);		//重设文件创建掩模
	return;
}

void restore_old_socket()
{
	if (pmt_socket_backuped != 0)
	{
		unlink(PMT_SERVER_SOCKET);
		link(PMT_SERVER_SOCKET_BACKUP, PMT_SERVER_SOCKET);
		unlink(PMT_SERVER_SOCKET_BACKUP);
		pmt_socket_backuped = 0;
	}
}

void close_new_pmt_socket()
{
	if (newpmtfd > 0)
	{
		int i;
		close(newpmtfd);
		newpmtfd = -1;
		for (i = 0; i < SIZE_OF_CONNECTIONS_TABLE; i++)
		{
			if (client_fd[i] != -1)
			{
				close(client_fd[i]);
				client_fd[i] = -1;
			}
		}
		restore_old_socket();
	}
}

static void on_termination(int signo)
{
	restore_old_socket();
	exit(0);
}

void print_copyright()
{
	printf("\n");
	printf("/============================================================\\\n");
	printf("| PMT Table Transformer                                      |\n");
	printf("| Version 1.5.1 Built 07/21/2009                             |\n");
	printf("| Copyright (c) 2009 Atmosphere Group                        |\n");
	printf("|                                                            |\n");
	printf("| This program is another implementation of the method which |\n");
	printf("| findout by ptjiang to let CCcam to process multi ecm.      |\n");
	printf("|                                                            |\n");
	printf("| Creator: linking_dai                                       |\n");
	printf("| Contact: linking_dai@163.com                               |\n");
	printf("| Forum:   http://www.xltvrobbs.net                          |\n");
	printf("\\============================================================/\n");
	printf("\n\n");
}

int main(int argc, char **argv)
{
	print_copyright();

	int run_as_daemon = 1, use_pmt_client_mode = -1;
	int i;
	for (i = 1; i < argc; i++)
	{
		if (strcmp(argv[i], "-d") == 0)
			run_as_daemon = 0;
		else if (strcmp(argv[i], "-s") == 0)
			use_pmt_client_mode = 0;
		else if (strcmp(argv[i], "-c") == 0)
			use_pmt_client_mode = 1;
		else if (strcmp(argv[i], "-l") == 0)
		{
			if (i < argc - 1)
			{
				delay = atoi(argv[i + 1]);
				i++;
			}
		}
	}

	if (use_pmt_client_mode < 0)
	{
		if (access("/bin/enigma", F_OK) == 0)
			use_pmt_client_mode = 0;
		else
			use_pmt_client_mode = 1;
	}

	if (run_as_daemon)
		init_daemon();

	int pmtfd = -1;
	struct sockaddr_un saddr;
	unsigned char pmtbuf[1024];
	int pmtlen;
	int new_pmt_server_threads_created = 0;
	int pmt_distribute_thread_created = 0;

	signal(SIGINT, on_termination);
	signal(SIGTERM, on_termination);

	pthread_mutex_init(&pmt_distribute_lock, NULL);
	pthread_cond_init(&new_pmt_arrived, NULL);

	memset(client_fd, -1, sizeof(client_fd));

	int error_count = 0;

	while (1)
	{
		if (!use_pmt_client_mode)
		{
			printf("Trying connect to PMT listen socket...\n");

			/* socket init */
			saddr.sun_family = AF_UNIX;
			if (pmt_socket_backuped != 0)
				strncpy(saddr.sun_path, PMT_SERVER_SOCKET_BACKUP, 107);
			else
				strncpy(saddr.sun_path, PMT_SERVER_SOCKET, 107);
			saddr.sun_path[107] = '\0';
			if (pmtfd == -1)
			{
				if ((pmtfd = socket(AF_UNIX, SOCK_STREAM, 0)) < 0)
				{
					printf("socket error (errno=%d)\n", errno);
					perror("");
					return 1;
				}
			}
			if (connect(pmtfd, (struct sockaddr *)&saddr, sizeof(saddr)) < 0) {
				printf("socket connect error (errno=%d)\n", errno);
				perror("");
				usleep(200000);
	
				error_count++;
				if (error_count >= 5) {
					close_new_pmt_socket();
					new_pmt_server_threads_created = 0;
					error_count = 0;
				}
				continue;
			}
	
			printf("Connected.\n");
	
			if (pmt_socket_backuped == 0)
			{
				unlink(PMT_SERVER_SOCKET_BACKUP);
				link(PMT_SERVER_SOCKET, PMT_SERVER_SOCKET_BACKUP);
				pmt_socket_backuped = 1;
			}
		}
		else
		{
			printf("Trying start original PMT listen socket...\n");

			/* socket init */
			saddr.sun_family = AF_UNIX;
			strncpy(saddr.sun_path, PMT_LISTEN_SOCKET, 107);
			saddr.sun_path[107] = '\0';
			unlink(saddr.sun_path);
			if (pmtfd == -1)
			{
				if ((pmtfd = socket(AF_UNIX, SOCK_STREAM, 0)) < 0)
				{
					printf("socket error (errno=%d)\n", errno);
					perror("");
					return 1;
				}
			}
			if (bind(pmtfd, (struct sockaddr *)&saddr,
				sizeof(saddr.sun_family)+strlen(saddr.sun_path) ) < 0) {
				printf("socket bind error (errno=%d)\n", errno);
				perror("");
				sleep(1);
				continue;
			}
			if (listen(pmtfd, 5) < 0) {
				printf("socket listen error (errno=%d)\n", errno);
				perror("");
				close(pmtfd);
				pmtfd = -1;
				sleep(1);
				continue;
			}
		}

		if (new_pmt_server_threads_created == 0)
		{
			pthread_t new_pmt_server_thread;
			pthread_attr_t new_pmt_server_thread_attr;
			pthread_attr_init (&new_pmt_server_thread_attr);
			pthread_attr_setdetachstate (&new_pmt_server_thread_attr, PTHREAD_CREATE_DETACHED);
			if(pthread_create(&new_pmt_server_thread, &new_pmt_server_thread_attr, new_pmt_server_proc, NULL) < 0)
			{
				printf("failed to create new pmt server thread (errno=%d).\n", errno);
				perror("");
				return 5;
			}
			new_pmt_server_threads_created = 1;
		}

		if (pmt_distribute_thread_created == 0)
		{
			pthread_t pmt_distribute_thread;
			pthread_attr_t pmt_distribute_thread_attr;
			pthread_attr_init (&pmt_distribute_thread_attr);
			pthread_attr_setdetachstate (&pmt_distribute_thread_attr, PTHREAD_CREATE_DETACHED);
			if(pthread_create(&pmt_distribute_thread, &pmt_distribute_thread_attr, pmt_distribute_proc, NULL) < 0)
			{
				printf("failed to create pmt distribute thread (errno=%d).\n", errno);
				perror("");
				return 5;
			}
			pmt_distribute_thread_created = 1;
		}

		printf("Waiting for PMTs...\n");

		if (!use_pmt_client_mode)
		{
			while (1)
			{
				pmtlen = recv(pmtfd, pmtbuf, sizeof(pmtbuf), 0);
				if (pmtlen < 0)
				{
					close(pmtfd);
					pmtfd = -1;
					printf("recv error (errno=%d)\n", errno);
					perror("");
					break;
				}
				else if (pmtlen == 0)
				{
					close(pmtfd);
					pmtfd = -1;
					printf("server connection closed.\n");
					break;
				}
	
				utils_dump("Original PMT: ", pmtbuf, pmtlen);
				pthread_mutex_lock(&pmt_distribute_lock);
				pmt_len = convert_pmt(pmtbuf, pmtlen, pmt_buffer);
				pthread_cond_signal(&new_pmt_arrived);
				pthread_mutex_unlock(&pmt_distribute_lock);
				utils_dump("New PMT: ", pmt_buffer, pmt_len);
			}
		}
		else
		{
			while (1)
			{
				fd_set readfds;
				FD_ZERO(&readfds);
				FD_SET(pmtfd, &readfds);
	
				struct timeval timeout;
				timeout.tv_sec = 5;
				timeout.tv_usec = 0;
		
				int nRetCode = select(pmtfd + 1, &readfds, 0, 0, &timeout);
				if(nRetCode > 0 /*&& FD_ISSET(pmtfd, &readfds)*/)
				{   
					struct sockaddr_in client_addr;
					socklen_t length;
					int clientfd;
		
					length = sizeof(client_addr);
					clientfd = accept(pmtfd,(struct sockaddr*)&client_addr,&length);
					if (clientfd < 0)
					{
						close(pmtfd);
						pmtfd = -1;
						printf("original PMT server accept error (errno=%d)\n", errno);
						perror("");
						continue;
					}
	
					while (1)
					{
						pmtlen = recv(clientfd, pmtbuf, sizeof(pmtbuf), 0);
						if (pmtlen < 0)
						{
							close(clientfd);
							printf("recv error (errno=%d)\n", errno);
							perror("");
							break;
						}
						else if (pmtlen == 0)
						{
							close(clientfd);
							printf("original PMT connection closed.\n");
							break;
						}
		
						utils_dump("Original PMT: ", pmtbuf, pmtlen);
						pthread_mutex_lock(&pmt_distribute_lock);
						pmt_len = convert_pmt(pmtbuf, pmtlen, pmt_buffer);
						pthread_cond_signal(&new_pmt_arrived);
						pthread_mutex_unlock(&pmt_distribute_lock);
						utils_dump("New PMT: ", pmt_buffer, pmt_len);
					}
				}
				else
				{
					close(pmtfd);
					pmtfd = -1;
					break;
				}
			}
		}
	}

	return 0;
}

